{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(logscale) pressure histo\t(entries=729)\n",
      "2.16                   ▉ \n",
      "   |                  ▉▉ \n",
      "   |                 ▉▉▉▉\n",
      "   |           ▉ ▉ ▉▉▉▉▉▉\n",
      "   |        ▉▉ ▉▉▉▉▉▉▉▉▉▉\n",
      "   |       ▉▉▉ ▉▉▉▉▉▉▉▉▉▉\n",
      "   |  ▉ ▉▉ ▉▉▉▉▉▉▉▉▉▉▉▉▉▉\n",
      "   | ▉▉▉▉▉ ▉▉▉▉▉▉▉▉▉▉▉▉▉▉\n",
      "   | ▉▉▉▉▉▉▉▉▉▉▉▉▉▉▉▉▉▉▉▉\n",
      "   | ▉▉▉▉▉▉▉▉▉▉▉▉▉▉▉▉▉▉▉▉\n",
      "-6.52....................0.28\n",
      " \n",
      "addPointScalars is DEPRECATED\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "921062f94e2c40e49c9cc134b1107d2b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Plot(antialias=True, axes=['x', 'y', 'z'], background_color=16777215, fps_meter=False, grid=[-1, -1, -1, 1, 1,…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\"\"\"\n",
    "Stokes equations with an iterative solver.\n",
    "\"\"\"\n",
    "# https://fenicsproject.org/docs/dolfin/2018.1.0/python/demos/\n",
    "#  stokes-iterative/demo_stokes-iterative.py.html\n",
    "from dolfin import *\n",
    "\n",
    "mesh = UnitCubeMesh(8, 8, 8)\n",
    "\n",
    "# Build function space\n",
    "P2 = VectorElement(\"Lagrange\", mesh.ufl_cell(), 2)\n",
    "P1 = FiniteElement(\"Lagrange\", mesh.ufl_cell(), 1)\n",
    "TH = P2 * P1\n",
    "W = FunctionSpace(mesh, TH)\n",
    "\n",
    "# Boundaries\n",
    "def right(x, on_boundary):\n",
    "    return x[0] > (1.0 - DOLFIN_EPS)\n",
    "\n",
    "def left(x, on_boundary):\n",
    "    return x[0] < DOLFIN_EPS\n",
    "\n",
    "def top_bottom(x, on_boundary):\n",
    "    return x[1] > 1.0 - DOLFIN_EPS or x[1] < DOLFIN_EPS\n",
    "\n",
    "# No-slip boundary condition for velocity\n",
    "noslip = Constant((0.0, 0.0, 0.0))\n",
    "bc0 = DirichletBC(W.sub(0), noslip, top_bottom)\n",
    "\n",
    "# Inflow boundary condition for velocity\n",
    "inflow = Expression((\"-sin(x[1]*pi)\", \"0.0\", \"0.0\"), degree=2)\n",
    "bc1 = DirichletBC(W.sub(0), inflow, right)\n",
    "\n",
    "# Define variational problem\n",
    "(u, p) = TrialFunctions(W)\n",
    "(v, q) = TestFunctions(W)\n",
    "f = Constant((0.0, 0.0, 0.0))\n",
    "a = inner(grad(u), grad(v)) * dx + div(v) * p * dx + q * div(u) * dx\n",
    "L = inner(f, v) * dx\n",
    "\n",
    "# Form for use in constructing preconditioner matrix\n",
    "b = inner(grad(u), grad(v)) * dx + p * q * dx\n",
    "\n",
    "# Assemble system\n",
    "A, bb = assemble_system(a, L, [bc0, bc1])\n",
    "\n",
    "# Assemble preconditioner system\n",
    "P, btmp = assemble_system(b, L, [bc0, bc1])\n",
    "\n",
    "# Create Krylov solver and AMG preconditioner\n",
    "if has_krylov_solver_method(\"minres\"):\n",
    "    krylov_method = \"minres\"\n",
    "elif has_krylov_solver_method(\"tfqmr\"):\n",
    "    krylov_method = \"tfqmr\"\n",
    "solver = KrylovSolver(krylov_method, \"amg\")\n",
    "\n",
    "# Associate operator (A) and preconditioner matrix (P)\n",
    "solver.set_operators(A, P)\n",
    "\n",
    "# Solve\n",
    "U = Function(W)\n",
    "solver.solve(U.vector(), bb)\n",
    "\n",
    "# Get sub-functions\n",
    "u, p = U.split()\n",
    "pressures = p.compute_vertex_values(mesh)\n",
    "\n",
    "\n",
    "#################################################### vtkplotter\n",
    "from vtkplotter.dolfin import plot, printHistogram\n",
    "\n",
    "printHistogram(pressures, title='pressure histogram', logscale=True)\n",
    "\n",
    "plot(u, wireframe=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
